from sqlalchemy import Engine, MetaData, Table, select, case, func, and_, or_
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import Session, sessionmaker, aliased

import constants
from data_model import ThreadTraceListData, LaunchTraceData


def query_thread_trace_tree_list(db_engine: Engine, process_id: int, thread_id: int, start_time: int, end_time: int):
    metadata = MetaData()
    metadata.reflect(bind=db_engine)
    Session = sessionmaker(bind=db_engine)
    session = Session()

    callstack = Table('callstack', metadata, autoload_with=db_engine)
    thread = Table('thread', metadata, autoload_with=db_engine)
    process = Table('process', metadata, autoload_with=db_engine)

    A = aliased(thread, name='A')
    C = aliased(callstack, name='C')
    IP = aliased(process, name='IP')

    query = (
        select(
            C.c.name,
            C.c.depth,
            A.c.tid,
            C.c.ts.label('startTime'),
            func.ifnull(A.c.name, 'Thread').label('threadName'),
            case((C.c.dur.is_(None), end_time - C.c.ts), else_=C.c.dur).label('duration')
        )
        .select_from(A)
        .join(C, C.c.callid == A.c.id)
        .join(IP, IP.c.ipid == A.c.ipid)
        .filter(
            C.c.cookie.is_(None),
            IP.c.pid == process_id,
            A.c.tid == thread_id,
            C.c.ts >= 0,
            and_(
                C.c.ts + case((C.c.dur.is_(None), end_time - C.c.ts), else_=C.c.dur) > start_time,
                C.c.ts < end_time
            )
        )
        .order_by(A.c.tid, C.c.ts)
    )

    result = session.execute(query)

    trace_tree_list = []
    for each in result.fetchall():
        data = ThreadTraceListData(each[0], each[1], each[2], each[3], each[4], each[5])
        trace_tree_list.append(data)
    return trace_tree_list


def query_launch_trace_list(db_engine: Engine):
    task_name_list = [
        constants.TRACE_START_ABILITY,
        constants.TRACE_START_ABILITY_LOCKED,
        constants.TRACE_LOAD_ABILITY,
        constants.TRACE_ATTACH_APPLICATION,
        constants.TRACE_LAUNCH_ABILITY,
        constants.TRACE_FOREGROUND_ABILITY,
        constants.TRACE_RECEIVE_VSYNC,
        constants.TRACE_APP_RENDERING,
        constants.TRACE_RS_RENDERING
    ]

    metadata = MetaData()
    metadata.reflect(bind=db_engine)
    Session = sessionmaker(bind=db_engine)
    session = Session()

    callstack = Table('callstack', metadata, autoload_with=db_engine)
    thread = Table('thread', metadata, autoload_with=db_engine)

    conditions = [callstack.c.cookie.is_(None)]
    like_conditions = or_(
        *[callstack.c.name.like(f"%{name}%") for name in task_name_list]
    )
    combined_conditions = and_(*conditions, like_conditions)

    query = (
        select(
            thread.c.tid.label('threadId'),
            case((thread.c.name.is_(None), 'Thread'), else_=thread.c.name).label('threadName'),
            callstack.c.ts.label('startTime'),
            case((callstack.c.dur.is_(None), -1), else_=callstack.c.dur).label('duration'),
            callstack.c.name
        ).select_from(
            callstack.join(
                thread,
                callstack.c.callid == thread.c.id
            )
        )
        .where(combined_conditions)
        .order_by(callstack.c.ts)
    )

    result = session.execute(query)

    launch_trace_list = []
    for each in result.fetchall():
        data = LaunchTraceData(each[0], each[1], each[2], each[3], each[4])
        launch_trace_list.append(data)
    return launch_trace_list


def get_trace_frame_slice_duration(db_engine, trace_name: str, thread_id: int, start_time: int):
    metadata = MetaData()
    metadata.reflect(bind=db_engine)
    Session = sessionmaker(bind=db_engine)
    session = Session()
    frame_slice = Table('frame_slice', metadata, autoload_with=db_engine)
    callstack = Table('callstack', metadata, autoload_with=db_engine)
    thread = Table('thread', metadata, autoload_with=db_engine)

    query = (
        select(frame_slice.c.dur)
        .select_from(frame_slice.join(callstack, frame_slice.c.callstack_id == callstack.c.id)
                     .join(thread, callstack.c.callid == thread.c.id))
        .where(
            callstack.c.ts == start_time,
            callstack.c.name == trace_name,
            frame_slice.c.type == 0,
            frame_slice.c.flag.in_([0, 1]),
            thread.c.tid == thread_id
        )
        .limit(1)
    )
    return session.execute(query).scalar_one_or_none()